# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/2 19:53 
@Author  : ZhangShenao 
@File    : 3.并行调用边.py 
@Desc    : 并行调用边

在LangGraph中,一个节点可以同时连接多条边
被连接到的所有节点全部都会并行执行
直到再次关联到一起,或者图运行结束

可以把状态看作是整个图架构的全局变量
每个节点的执行,都是一次全局变量的修改操作
直到执行到END节点,最终输出全局变量
"""
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.constants import START
from langgraph.graph import MessagesState, StateGraph


# 定义llm节点
def llm_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
    msg = AIMessage(content="你好，我是智能聊天助手，请问有什么可以帮助您的？")
    return {"messages": [msg]}


# 定义两个并行节点
def parallel1(state: MessagesState, config: RunnableConfig) -> MessagesState:
    msg = HumanMessage(content="执行并行节点1")
    return {"messages": [msg]}


def parallel2(state: MessagesState, config: RunnableConfig) -> MessagesState:
    msg = HumanMessage(content="执行并行节点2")
    return {"messages": [msg]}


def chat_end(state: MessagesState, config: RunnableConfig) -> MessagesState:
    msg = HumanMessage(content="聊天结束")
    return {"messages": [msg]}


# 初始化图架构,指定状态为MessagesState类型
graph_builder = StateGraph(MessagesState)

# 添加节点
graph_builder.add_node(node="llm", action=llm_node)
graph_builder.add_node(node="parallel1", action=parallel1)
graph_builder.add_node(node="parallel2", action=parallel2)
graph_builder.add_node(node="chat_end", action=chat_end)

# 添加边
graph_builder.add_edge(start_key=START, end_key="llm")
graph_builder.add_edge(start_key="llm", end_key="parallel1")
graph_builder.add_edge(start_key="llm", end_key="parallel2")
graph_builder.add_edge(start_key="parallel1", end_key="chat_end")
graph_builder.add_edge(start_key="parallel2", end_key="chat_end")

# 编译图架构,生成Runnable可运行组件
graph = graph_builder.compile()

# 定义初始状态
init_state: MessagesState = {"messages": [HumanMessage(content="你好，我是zsa，我喜欢写代码和打篮球")]}

# 执行图,获取最终状态
final_state = graph.invoke(init_state)

for msg in final_state["messages"]:
    print(f"消息内容: {msg.content}")
